#include <mpi.h>
#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <math.h>
#include "common.h"

// MPI version of our serial program

#define IDX(i,j,NB) ((i)*(NB)+(j))
#define OUTBOUND(i,j,NB) ((i) < 0 || (i) >= (NB) || (j) < 0 || (j) >= (NB))

inline int getBinID(particle_t & particle, int NB)
{
    return floor(particle.x / cutoff) * NB + floor(particle.y / cutoff); 
}

extern double size;

void clear_bins(int NB2, int * binMax)
{
    // should we do this
    for (int k = 0; k < NB2; k ++)
    {
        binMax[k] = 0;
    }
    // or just this?
    //memset(binMax,0, sizeof(int) * NB2);
}

// note: bins is passed by reference, not by value.
void assign_bins(int NB, int n, particle_t *particles, int **bins, int *binMax)
{
    clear_bins(NB*NB, binMax);
    for (int k = 0; k < n; k++)
    {
        int binid = getBinID(particles[k],NB);
        bins[binid][binMax[binid]++] = k;
    }
}

void apply_force_binned(int NB, int ** bins, int * binMax, particle_t * particles, particle_t &particle)
{
    int dstbinid = getBinID(particle, NB);
    int i = dstbinid / NB;
    int j = dstbinid % NB;
    int istart = max(i-1,0), iend = min(i+2,NB),
        jstart = max(j-1,0), jend = min(j+2,NB);
    for (int _i = istart; _i < iend; _i ++)
    { 
        for (int _j = jstart; _j < jend; _j ++)
        {
            int * neighborbin = bins[IDX(_i,_j,NB)];
            int pnum = binMax[IDX(_i,_j,NB)];
            for (int src = 0; src < pnum; src ++)
            { 
	      apply_force(particle, particles[neighborbin[src]]);
            } 
        }
    }
}


int main( int argc, char **argv )
{   
    if( find_option( argc, argv, "-h" ) >= 0 )
    {
        printf( "Options:\n" );
        printf( "-h to see this help\n" );
        printf( "-n <int> to set the number of particles\n" );
        printf( "-o <filename> to specify the output file name\n" );
        return 0;
    }

    // set up MPI
    int n_proc, rank;
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &n_proc);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    int n = read_int( argc, argv, "-n", 1000 );
    char *savename = read_string( argc, argv, "-o", NULL );

    FILE *fsave = savename && rank == 0 ? fopen( savename, "w" ) : NULL;
    particle_t *particles = (particle_t*) malloc( n * sizeof(particle_t) );

    set_size( n );
    int NB = floor( size / cutoff) + 1; // NB is the number of bins per edge
 
    int P_PER_NB = 100; // should probably depend on size...
    // magical number: we assume that this is enough for each bin 

    int * bins[NB*NB];
    int binMax[NB*NB];
    for (int i = 0; i < NB*NB; i++) {
      bins[i] = new int[P_PER_NB];
      binMax[i] = 0;
    }

    MPI_Datatype PARTICLE;
    MPI_Type_contiguous( 6, MPI_DOUBLE, &PARTICLE );
    MPI_Type_commit( &PARTICLE );

    int particles_per_proc = (n + n_proc - 1) / n_proc;
    int *partition_offsets = (int*) malloc( (n_proc+1) * sizeof(int) );
    for( int i = 0; i < n_proc+1; i++ )
      partition_offsets[i] = min( i * particles_per_proc, n );
    
    int *partition_sizes = (int*) malloc( n_proc * sizeof(int) );
    for( int i = 0; i < n_proc; i++ )
      partition_sizes[i] = partition_offsets[i+1] - partition_offsets[i];

    //  allocate storage for local partition
    int nlocal = partition_sizes[rank];
    particle_t *local = (particle_t*) malloc( nlocal * sizeof(particle_t) );
    
    //  initialize and distribute the particles (that's fine to leave it unoptimized)
    // do this on only one processor, or else the starting locations will
    // be different on all of them!
    if (rank == 0) {
      init_particles( n, particles );
    }

    MPI_Scatterv( particles, partition_sizes, partition_offsets, PARTICLE, local, nlocal, PARTICLE, 0, MPI_COMM_WORLD );

    // simulate a number of time steps
    double simulation_time = read_timer( );

    for( int step = 0; step < NSTEPS; step++ )
    {   
      //printf("I am processor %i, now on step %i-0 \n", rank, step);

      // communicate full set of particles to everyone
      MPI_Allgatherv(local, nlocal, PARTICLE, particles, partition_sizes, partition_offsets, PARTICLE, MPI_COMM_WORLD);
      //      printf("I am processor %i, now on step %i-1 \n", rank, step);

      // everyone updates bins                                            
      assign_bins(NB, n, particles, bins, binMax);
      //printf("I am processor %i, now on step %i-2 \n", rank, step);

      //  save if necessary
      if( fsave && (step%SAVEFREQ) == 0 ) {
	save( fsave, n, particles );
      }
      //printf("I am processor %i, now on step %i-3 \n", rank, step);
      
      // everyone computes all forces on only THEIR particles
      for (int i = 0; i < nlocal; i++)
        {
            local[i].ax = local[i].ay = 0;
            apply_force_binned(NB, bins, binMax, particles, local[i]);
        }

      // everyone moves only THEIR particles
      for( int i = 0; i < nlocal; i++ )
        {
	  move( local[i] );
        }

    }
     
    simulation_time = read_timer( ) - simulation_time;

    if (rank == 0) {
      printf( "n = %d, simulation time = %g seconds\n", n, simulation_time );
    }

    free(partition_offsets);
    free(partition_sizes);
    free(local);
    free(particles);

    for (int i = 0; i < NB*NB; i ++) {
        delete bins[i];
    }

    if( fsave )
        fclose( fsave );

    MPI_Finalize();

    return 0;
}

